import cv2
import os  # 文件的读取遍历

import numpy as np


# recognizer.train()  # 训练特征模型
# recognizer.read()  # 加载特征模型
# recognizer.predict()  # 对图像进行预测

def face_detection(img):
    classifier = cv2.CascadeClassifier("./haarcascades/haarcascade_frontalface_alt2.xml")
    faces = classifier.detectMultiScale(
        image=img,
        scaleFactor=1.2,
        minNeighbors=5,
        flags=0,
        minSize=[50, 50],
        maxSize=[200, 200]
    )

    return faces[0]


def get_frames_ids():
    files = os.listdir("./users")
    frames = []
    ids = []
    for filename in files:
        img = cv2.imread(os.path.join("./users", filename))
        x, y, w, h = face_detection(img)
        id = int(filename.split(".")[0])
        # 将图像设置为单通道的灰度图
        gray_face = cv2.cvtColor(img[y:y + h, x:x + w], cv2.COLOR_BGR2GRAY)
        frames.append(gray_face)
        ids.append(id)
    return frames, ids


if __name__ == "__main__":
    frames, ids = get_frames_ids()
    print(frames,ids)
    # for frame in frames:
    #     print(frame.shape)
    # TODO：此处为 LBPHFaceRecognizer_create
    recognizer = cv2.face.LBPHFaceRecognizer_create()
    recognizer.train(frames, np.array(ids))
    recognizer.write('./data/test.yml')

